#include "StunAgent.h"
#include "Poco/ByteOrder.h"
#include "Poco/Net/IPAddress.h"

using Poco::ByteOrder;
using Poco::Net::IPAddress;

namespace cm_nat
{
    StunAgent::StunAgent( ByteBufferPool* p_pool, BaseTimer* p_timer )
        :m_bufferPool(p_pool), m_timer(p_timer),
        m_keepCheckFlag(false),m_connectFlag(false),
        m_checkFailFlag(true),
        m_checkFailCount(0)
    {

    }

    StunAgent::~StunAgent()
    {

    }

    int StunAgent::setup()
    {
        return CM_RES_SUCCESS;
    }

    int StunAgent::teardown()
    {
        return CM_RES_SUCCESS;
    }

    void StunAgent::registerListener(StunAgentListener* p_listener)
    {
        m_listener = p_listener;
    }

    int StunAgent::startCheck( SocketAddress remote )
    {
        m_remoteAddress = remote;
        m_connectFlag = false;
        m_lastTid = generateTid();
        sendRequest( m_remoteAddress, m_lastTid );
        startTimer(5000, TIMER_START_CHECK);
        
        return CM_RES_SUCCESS;
    }

    int StunAgent::keepCheck()
    {
        m_keepCheckFlag = true;
        m_lastTid = generateTid();
        sendRequest( m_remoteAddress , m_lastTid);
        startTimer(5000, TIMER_KEEP_CHECK);

        //poco_debug(g_log_nat, "StunAgent::keepCheck()");

        return CM_RES_SUCCESS;
    }

    int StunAgent::stopCheck()
    {
        m_keepCheckFlag = false;

        return CM_RES_SUCCESS;
    }

    SocketAddress StunAgent::getReflexAddress()
    {
        return m_reflexAddress;
    }
    
    bool StunAgent::filterStunPacket( ByteBuffer* p_data )
    {
        if( detectMessage(p_data) ){
            processMessage(p_data);

            SocketAddress* p_sa = (SocketAddress*)p_data->getData();
            delete p_sa;
            p_data->decreaseRef();

            return true;
        }else{
            return false;
        }
    }

    int StunAgent::sendRequest(SocketAddress remote, std::string& tid)
    {
        StunMessage smsg;
        smsg.SetType(STUN_BINDING_REQUEST);
        smsg.SetTransactionID(tid);

        ByteBuffer* p_data = NULL;
        p_data = m_bufferPool->get(CM_NO_WAIT);
        
        if( !p_data ){
            return CM_RES_FAIL;
        }

        p_data->increaseRef();
        smsg.Write( p_data );

        sendTo( p_data, remote);

        return CM_RES_SUCCESS;
    }

    bool StunAgent::detectMessage(ByteBuffer* data)
    {
        if(data->Length()<20){
            return false;
        }

        const char* p_data = data->Data();
        char byte = p_data[0];

        if( byte & 0xc0 ){
            return false;
        }
        //stun message length is 4 aligns
        byte = p_data[3];
        if( byte & 0x03 ){
            return false;
        }

        if( (p_data[4]&0x21) && (p_data[5]&0x12) && (p_data[6]&0xa4) && ( p_data[7]&0x42) ){
            return true;
        }
        else{
            return false;
        }
    }

    int StunAgent::processMessage( ByteBuffer* data )
    {
        bool res = false;
        StunMessage smsg;
        res = smsg.Read( data );

        const SocketAddress* from = (SocketAddress*)data->getData();

        switch( smsg.type() ){

        case STUN_BINDING_REQUEST:
            {
                processRequest( &smsg, *from );
            }
            break;
        case STUN_BINDING_RESPONSE:
            {
                processResponse( &smsg, *from );
            }
            break;
        case STUN_BINDING_ERROR_RESPONSE:
            {
                processErrorResponse( &smsg, *from );
            }
            break;
        default:
            break;
        }

        return true;

    }

    int StunAgent::processResponse( StunMessage* p_sm,const SocketAddress& from  )
    {
        const StunAddressAttribute* p_ma = p_sm->GetAddress(STUN_ATTR_MAPPED_ADDRESS);
        if( !p_ma ){
            return CM_RES_FAIL;
        }
        //get reflex address
        struct in_addr  ip;
        ip.s_addr = ByteOrder::toNetwork( (UInt32)p_ma->ip() );
        SocketAddress sa( IPAddress(&ip, sizeof(in_addr)), p_ma->port() );

        if( p_sm->transaction_id()==m_lastTid ){

            if( !m_connectFlag )
            {
                m_reflexAddress = sa;
                m_connectFlag = true;
                if( m_listener ){
                    m_listener->onCheckResult( this, StunAgentListener::CHECK_RES_CHECK_SUCCESS );
                }
            }
            else
            {
                if( m_reflexAddress.host()!=sa.host() || m_reflexAddress.port()!=sa.port() )
                {
                    m_reflexAddress = sa;
                    if( m_listener ){
                        m_listener->onCheckResult( this, StunAgentListener::CHECK_RES_REFLEX_UPDATE );
                    }
                }
            }

            {
                FastMutex::ScopedLock sl(m_mutex);
                m_checkFailFlag = false;
            }

        }

        return CM_RES_SUCCESS;
    }

    int StunAgent::processErrorResponse( StunMessage* p_sm,const SocketAddress& from )
    {
        return CM_RES_SUCCESS;
    }

    int StunAgent::processRequest( StunMessage* p_sm,const SocketAddress& from )
    {
        StunMessage response;
        response.SetType(STUN_BINDING_RESPONSE);
        response.SetTransactionID( p_sm->transaction_id() );

        StunAddressAttribute* addr_attr = 
            StunAttribute::CreateAddress(STUN_ATTR_MAPPED_ADDRESS);

        addr_attr->SetFamily(1);
        addr_attr->SetPort( from.port() );

        IPAddress from_ip = from.host();
        const struct in_addr* ip= (struct in_addr*)from_ip.addr();
        addr_attr->SetIP( ByteOrder::fromNetwork( (UInt32)ip->s_addr ) );
        response.AddAttribute(addr_attr);

        ByteBuffer* p_data = NULL;
        if(m_bufferPool){
            p_data = m_bufferPool->get(0);
        }else{
            return false;
        }
        if( !p_data ){
            poco_debug(g_log_nat,"StunAgent::processIncomingRequest get buffer fail");
            return false;
        }
        p_data->increaseRef();
        response.Write( p_data );

        sendTo( p_data, from);

        return CM_RES_SUCCESS;
    }

    int StunAgent::sendTo(ByteBuffer* data,const SocketAddress& sa)
    {
        poco_assert( m_listener != NULL );
        m_listener->onToSend( this, data, sa );

        return CM_RES_SUCCESS;
    }

    std::string StunAgent::generateTid()
    {
        std::string tid;
        tid.resize(16);

        //RFC5389 magic cookie
        tid[0] = 0x21;
        tid[1] = 0x12;
        tid[2] = 0xA4;
        tid[3] = 0x42;

        for(int i =4; i<16; i++){
            tid[i] = m_random.nextChar();
        }
        return tid;
    }

    void StunAgent::startTimer(long interval, UInt32 context)
    {
        m_timer->start(interval, context, this );
    }


    void StunAgent::onTimer(TimerNotification* p_tn)
    {
        if(! p_tn ){
            return;
        }
        //timer and onReceive() task may in different thread
        //need to protect some vars.
        FastMutex::ScopedLock sl(m_mutex);

        switch( p_tn->getContext() )
        {
        case TIMER_START_CHECK:
            {
                if( !m_connectFlag )
                {
                    if( m_listener ){
                        m_listener->onCheckResult(this, StunAgentListener::CHECK_RES_CHECK_FAIL);
                    }
                }
            }
            break;

        case TIMER_KEEP_CHECK:
            {
                if( m_checkFailFlag )
                {
                    m_checkFailCount++;

                    if( m_checkFailCount > KEEP_CHECK_FAIL_MAX ){
                        if( m_listener ){
                            m_listener->onCheckResult(this, StunAgentListener::CHECK_RES_KEEP_FAIL );
                        }    
                    }
                    else{
                        //try check again
                        keepCheck();
                    }
                }
                else
                {
                    //reset fail count
                    m_checkFailCount = 0;
                    keepCheck();
                }

                m_checkFailFlag = true;
            }
            break;

        default:
            break;
        }
    }


}
